package com.feidee.fd.sml.algorithm.component.ml.regression

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.regression.DecisionTreeRegressor

/**
  * @Author songhaicheng
  * @Date 2019/3/25 16:05
  * @Description
  * @Reviewer dongguosheng
  */
case class DTRegressionParam(
                              override val input_pt: String,
                              override val output_pt: String,
                              override val hive_table: String,
                              override val flow_time: String,
                              override val featuresCol: String,
                              override var labelCol: String,
                              override var predictionCol: String,
                              override val modelPath: String,
                              override val metrics: Array[String],
                              // 误差列，默认空
                              varianceCol: String,
                              // 决策树节点特征选择方式，仅支持 variance
                              impurity: String,
                              // 树最大深度，>= 0，默认 5
                              maxDepth: Int,
                              // 最大划分树，>= 2，默认 32
                              maxBins: Int,
                              // 每个节点最少实例数，>= 1，默认 1
                              minInstancesPerNode: Int,
                              // 检查点间隔，-1 或 >= 1,-1 时不做缓存，默认 10
                              checkpointInterval: Int,
                              // 最大内存 MB 单位，这个值越大，一次处理的节点划分就越多，>= 0，默认 256
                              maxMemoryInMB: Int,
                              // 随机种子，默认 12346
                              seed: Long,
                              // 最小信息增益，>= 0.0，默认 0.0
                              minInfoGain: Double,
                              // 是否缓存节点 id，缓存可以加速深层树的训练，默认 false
                              cacheNodeIds: Boolean
                            ) extends MLParam {
  def this() = this(null, null, null, null, "features", "label", "prediction", null, new Array[String](0),
    null, "variance", 5, 32, 1, 10, 256, 123456L, 0.0, false)

  override def verify(): Unit = {
    super.verify()
    val impurities = Array("variance")
    require(impurities.contains(impurity.toLowerCase), s"param impurity only accepts " +
      s"${impurities.mkString("[", ", ", "]")}, but has $impurity")
    require(maxDepth >= 0, "param maxDepth can't be negative")
    require(maxBins >= 2, "param maxBins must be not less than 2")
    require(minInstancesPerNode >= 1, "param numTrees must be not less than 1")
    require(checkpointInterval == -1 || checkpointInterval >= 1, "param checkpointInterval must be" +
      " equals to -1 or not less than 1")
    require(minInfoGain >= 0, "param minInfoGain can't be negative")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("varianceCol" -> varianceCol)
    map += ("impurity" -> impurity)
    map += ("maxDepth" -> maxDepth)
    map += ("maxBins" -> maxBins)
    map += ("minInstancesPerNode" -> minInstancesPerNode)
    map += ("checkpointInterval" -> checkpointInterval)
    map += ("maxMemoryInMB" -> maxMemoryInMB)
    map += ("seed" -> seed)
    map += ("minInfoGain" -> minInfoGain)
    map += ("cacheNodeIds" -> cacheNodeIds)
    map
  }

}


class DTRegressionComponent extends AbstractRegressionComponent[DTRegressionParam] {

  override def setUp(param: DTRegressionParam): PipelineStage = {
    val dtr = new DecisionTreeRegressor()
      .setLabelCol(param.labelCol)
      .setFeaturesCol(param.featuresCol)
      .setPredictionCol(param.predictionCol)
      .setImpurity(param.impurity)
      .setMaxDepth(param.maxDepth)
      .setMaxBins(param.maxBins)
      .setMinInstancesPerNode(param.minInstancesPerNode)
      .setCheckpointInterval(param.checkpointInterval)
      .setMaxMemoryInMB(param.maxMemoryInMB)
      .setSeed(param.seed)
      .setMinInfoGain(param.minInfoGain)
      .setCacheNodeIds(param.cacheNodeIds)

    if (tool.isNotNull(param.varianceCol)) {
      dtr.setVarianceCol(param.varianceCol)
    }

    dtr
  }

}

object DTRegressionComponent {
  def apply(paramStr: String): Unit = {
    new DTRegressionComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    DTRegressionComponent(args(0))
  }
}
